package com.eudon.ai.agent.etl.factory;

import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.generation.augmentation.QueryAugmenter;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;

/**
 * @author eudon
 * @description 检索增强器工厂，和ChatClient.prompt().advisors()搭配使用，增加知识库检索能力。
 * @date 2025/8/16
 * @email syd19970616@gmail.com
 * @address <a href="https://gitee.com/eudon">Gitee</a>
 */
public class RetrievalAugmentationAdvisorFactory {
    private static Advisor doCreate(VectorStore myPostgresVectorStore, Filter.Expression filter) {
        DocumentRetriever documentRetriever;
        if (filter == null) {
            documentRetriever = VectorStoreDocumentRetriever.builder()
                    .similarityThreshold(0.8)
                    .vectorStore(myPostgresVectorStore)
                    .build();

        } else {
            documentRetriever = VectorStoreDocumentRetriever.builder()
                    .filterExpression(filter)
                    .similarityThreshold(0.8)
                    .vectorStore(myPostgresVectorStore)
                    .build();
        }

        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(documentRetriever)
//                .documentJoiner()   //文档合并
//                .queryExpander()    //查询扩展器
//                .queryTransformers()    //查询转换器，官方提供了查询重写器、查询翻译器、查询压缩器三种实现，可以自行实现功能
                .queryAugmenter(DefaultQueryAugmentBuilder.build())   //查询增强器
                .order(9)    //排序，建议该拦截器顺序在日志拦截器前一个执行
                .build();
    }

    public static Advisor create(VectorStore myPostgresVectorStore, String status, String label) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        Filter.Expression expression = b.and(b.eq("status", status), b.eq("keywords", label)).build();
        return doCreate(myPostgresVectorStore, expression);
    }

    public static Advisor create(VectorStore myPostgresVectorStore, String status) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        Filter.Expression expression = b.eq("status", status).build();
        return doCreate(myPostgresVectorStore, expression);
    }

    public static Advisor create(VectorStore myPostgresVectorStore) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();
        return doCreate(myPostgresVectorStore, null);
    }

    public static class DefaultQueryAugmentBuilder {
        private static final PromptTemplate EMPTY_CONTEXT_PROMPT_TEMPLATE = new PromptTemplate("你应该输出下面的内容：\n" +
                "我是您的专属恋爱小助手，但是我只能回答恋爱相关的问题，您可以重新向我提问恋爱相关的问题哦～");

        public static QueryAugmenter build() {
            return ContextualQueryAugmenter.builder()
                    .allowEmptyContext(true)    //默认为false，设置为true时增强器会返回用户原本请求，空上下文模版功能会失效
                    .emptyContextPromptTemplate(EMPTY_CONTEXT_PROMPT_TEMPLATE)
                    .build();
        }
    }
}
